import numpy as np
import platgo as pg


class Sphere(pg.Problem):

    def __init__(self, D: int = 20) -> None:
        self.name = 'Sphere'
        self.type['single'], self.type['real'], self.type['expensive'] = [True] * 3
        self.M = 1
        self.D = D
        lb = [-100] * self.D
        ub = [100] * self.D
        self.borders = np.array([lb, ub])
        super().__init__()

    def cal_obj(self, pop: pg.Population) -> None:
        x = pop.decs
        pop.objv = np.array([np.sum(np.square(x), 1)]).T

    def get_optimal(self) -> np.ndarray:
        return np.array([[0]])

    def g_fun(self, pop: pg.Population):
        # 求梯度的函数
        gk = 2 * pop.decs
        return gk


if __name__ == '__main__':
    d = Sphere()
    pop = pg.Population(decs=np.random.uniform(0, 1, (10, 7)))
    print(d.borders)
    print(pop.decs)
    d.cal_obj(pop)  # 计算目标函数值
    print(pop.objv)
